#include <linux/module.h>
#include <linux/init.h> 
#include <linux/kernel.h>
#include <linux/skbuff.h>
#include <net/tcp.h>
#include <linux/netdevice.h>
#include <linux/netfilter.h>
#include <linux/netfilter_ipv4.h>
#include <linux/fs.h>
#include <linux/time.h>
#include <linux/rtc.h>
#include <linux/spinlock.h>
#include <linux/workqueue.h>
#include <linux/fs.h>
#include <linux/uaccess.h>
#include "myfirewall.h"


static struct nf_hook_ops nfhoLocalIn;  	    // 在数据路由后处理本机数据包的钩子
static struct nf_sockopt_ops nfhoSockopt;       // 内核与用户通信钩子

// 防火墙过滤规则
ban_status rules, recv;

// 状态检测Hash 表
ktime_t  hashTable[TABLE_SIZE] = {0};

// Hash 锁
char hashLock = 0;

static Log logs[LOG_NUM_MAX] = {0};

// 未写入日志文件的日志数
static int log_num = 0;

// 状态连接链表的表头 表尾
Connection connHead, connEnd;

// Hash函数
static unsigned get_hash(int k) 
{
	unsigned a, b, c = 4;
    a = b = 0x9e3779b9;
    a += k;
	a -= b; a -= c; a ^= (c >> 13); 
	b -= c; b -= a; b ^= (a << 8); 
	c -= a; c -= b; c ^= (b >> 13); 
	a -= b; a -= c; a ^= (c >> 12);  
	b -= c; b -= a; b ^= (a << 16); 
	c -= a; c -= b; c ^= (b >> 5); 
	a -= b; a -= c; a ^= (c >> 3);  
	b -= c; b -= a; b ^= (a << 10); 
	c -= a; c -= b; c ^= (b >> 15); 

    return c % TABLE_SIZE;
}

/**
 * 这个刚到达的数据包，是不是属于一个我们已经认识并且允许通过的连接
 * 检测该数据包是否已经建立连接
 *  */ 
bool check_conn(struct sk_buff *skb){

	struct iphdr *ip = ip_hdr(skb);
	unsigned int src_ip = ntohl(ip->saddr);
	unsigned int dst_ip = ntohl(ip->daddr);

	int src_port;
	int dst_port;
	int protocol;
	unsigned int scode;
	unsigned int pos;

	if(!skb) return true;

	if(ip->protocol == IPPROTO_TCP){
		struct tcphdr *tcp = tcp_hdr(skb);
		src_port = ntohs(tcp->source);
		dst_port = ntohs(tcp->dest);
		protocol = TCP;
	}else if(ip->protocol == IPPROTO_UDP){
		struct udphdr *udp = udp_hdr(skb);
		src_port = ntohs(udp->source);
		dst_port = ntohs(udp->dest);
		protocol = UDP;
	}else if(ip->protocol == IPPROTO_ICMP){
        // 对ICMP协议，端口被硬编码为30001
		src_port = 30001;
		dst_port = 30001;
		protocol = ICMP;
	}else{
		// 不记录状态
		return false;
	}

	// 3. 计算哈希编码和哈希表位置
    // 使用IP地址、端口和协议的异或组合作为哈希编码
	scode = src_ip ^ dst_ip ^ src_port ^ dst_port ^ protocol;
	pos = get_hash(scode);

	while(hashLock);
	hashLock = 1;

	if(hashTable[pos]){
		// 更新连接时间
		hashTable[pos] = CONNECT_TIME;
		// 开锁
		hashLock = 0;
		return true;
	}

	hashLock = 0;
	return false;
}

//获取系统时间
void get_time(char *time_buf, int len){
	ktime_t k_time;
	struct rtc_time tm;

	k_time = ktime_get_real();
	tm = rtc_ktime_to_tm(k_time);

	snprintf(time_buf, len, "%d-%d-%d %d:%d:%d",
	tm.tm_year + 1900, tm.tm_mon + 1, tm.tm_mday, tm.tm_hour + 8, tm.tm_min, tm.tm_sec);
}

ktime_t get_seconds(void)
{
    return ktime_get_real_seconds();
}

void convert_ip(unsigned int ip, char* ip_str, size_t size){
	snprintf(ip_str, size, "%u.%u.%u.%u",
	(ip & 0x000000ff) >> 0,
	(ip & 0x0000ff00) >> 8,
	(ip & 0x00ff0000) >> 16,
	(ip & 0xff000000) >> 24);
}

// 增加日志到日志数组
bool add_log(struct sk_buff* skb, char *rule_str){
	struct iphdr *iph = ip_hdr(skb);
	char time_buf[64];
	get_time(time_buf, sizeof(time_buf));

	if(log_num == LOG_NUM_MAX) {
		printk("The currently retained logs have reached the maximum value. \n");
		return false;
	}

	logs[log_num].src_ip = iph->saddr;
	logs[log_num].dst_ip = iph->daddr;

	if (iph->protocol == IPPROTO_TCP){
		struct tcphdr *tcph = tcp_hdr(skb);
		logs[log_num].src_port = ntohs(tcph->source);
		logs[log_num].dst_port = ntohs(tcph->dest);
		strcpy(logs[log_num].protocol, "TCP");
	}
	else if( iph->protocol == IPPROTO_UDP){
		struct udphdr *udph = udp_hdr(skb);
		logs[log_num].src_port = ntohs(udph->source);
		logs[log_num].dst_port = ntohs(udph->dest);
		strcpy(logs[log_num].protocol, "UDP");    
	}else if(iph->protocol == IPPROTO_ICMP){
		logs[log_num].src_port = 30001; 
		logs[log_num].dst_port = 30001;       
		strcpy(logs[log_num].protocol, "ICMP");	
	}else{
		logs[log_num].src_port = 0;
		logs[log_num].dst_port = 0;
		strcpy(logs[log_num].protocol, "ELSE");
	}

	memcpy(logs[log_num].curr_time, time_buf, sizeof(time_buf));
	memcpy(logs[log_num].filter_rule, rule_str, sizeof(logs[log_num].filter_rule));
	log_num++;

	return true;
}

// 更新状态连接表
void update_hashTable(struct sk_buff *skb){
	struct iphdr *iph = ip_hdr(skb);
	unsigned int src_ip = ntohl(iph->saddr);
	unsigned int dst_ip = ntohl(iph->daddr);

	int src_port;
	int dst_port;
	int protocol;
	unsigned int scode;
	unsigned int pos;
	Connection *conn_node;

	if(rules.connNum < CONN_NUM_MAX){
		// 区分协议类型
		if(iph->protocol == IPPROTO_TCP){
			struct tcphdr *tcp = tcp_hdr(skb);
			src_port = ntohs(tcp->source);
			dst_port = ntohs(tcp->dest);
			protocol = TCP; 
		}else if(iph->protocol == IPPROTO_UDP){
			struct udphdr *udp = udp_hdr(skb);
			src_port = ntohs(udp->source);
			dst_port = ntohs(udp->dest);
			protocol = UDP;
		}else if (iph->protocol == IPPROTO_ICMP) {
			src_port = 30001;
			dst_port = 30001;
			protocol = ICMP;
		}else {
			src_port = 0;
			dst_port = 0;
		}

		scode = src_ip ^ dst_ip ^ src_port ^ dst_port ^ protocol;
		pos = get_hash(scode);

		while(hashLock);
		hashLock = 1;

		hashTable[pos] = CONNECT_TIME;
		hashLock = 0;

		++rules.connNum;
		conn_node = (Connection *)kmalloc(sizeof(Connection), GFP_ATOMIC);
		conn_node->src_ip = src_ip;
		conn_node->dst_ip = dst_ip;
		conn_node->src_port = src_port;
		conn_node->dst_port = dst_port;
		conn_node->protocol = protocol;
		conn_node->index = pos;

		// 头插法
		conn_node->next = connHead.next;
		connHead.next = conn_node;
	}else{
		printk("The number of connections exceeds the maximum.\n");
	}
}

// 获取数据包端口号
void get_port(struct sk_buff *skb, unsigned short *src_port, unsigned short *dst_port){
	struct iphdr *iph = ip_hdr(skb);
	struct tcphdr *tcph;
	struct udphdr *udph;

	switch(iph->protocol){
		case IPPROTO_TCP:
			tcph = tcp_hdr(skb);
			*src_port = ntohs(tcph->source);
			*dst_port = ntohs(tcph->dest);
			break;
		case IPPROTO_UDP:
			udph = udp_hdr(skb);
			*src_port = ntohs(udph->source);
			*dst_port = ntohs(udph->dest);
			break;
		case IPPROTO_ICMP:
			*src_port = 30001;
			*dst_port = 30001;
			break;
		default:
			*src_port = 0;
			*dst_port = 0;
			break;
	}
}

// 判断是否是ping 包
bool is_PING(struct iphdr *iph){
	if(iph->protocol == IPPROTO_ICMP){
		return true;
	}
	return false;
}

// 判断是否是HTTP/HTTPS包
bool is_HTTP(struct iphdr *iph, struct tcphdr *tcph){
	if(iph->protocol == IPPROTO_TCP){
		if(tcph->dest == htons(80) || tcph->dest == htons(443) || tcph->dest == htons(8080)
		|| tcph->source == htons(80) || tcph->source == htons(443) || tcph->source == htons(8080)){
			return true;
		}
	}
	return false;
}

// 判断是否是Telnet包
bool is_TELNET(struct iphdr *iph, struct tcphdr *tcph){
	if(iph->protocol == IPPROTO_TCP && tcph->dest == htons(23)){
		return true;
	}
	return false;
}

/**
 * 检查到达的数据包，并作出裁决
 */
unsigned int hookLocalIn(void* priv, struct sk_buff* skb, const struct nf_hook_state* state){
	//  sk_buff是Linux 内核中用于管理网络数据包的核心数据结构

	// 从skb中提取IP协议头 header -> hdr
	struct iphdr *iph = ip_hdr(skb);
	struct tcphdr *tcph;

	//定义几个字符数组，用于存放转换成字符串格式的IP地址和当前时间
	char src_ip_str[16];
	char dst_ip_str[16];
	char time_buf[64];

	// 将二进制的IP地址和时间戳转换成人类可读的字符串，方便后续 printk 打印日志。
	convert_ip(iph->saddr, src_ip_str, sizeof(src_ip_str));
	convert_ip(iph->daddr, dst_ip_str, sizeof(dst_ip_str));
	get_time(time_buf, sizeof(time_buf));

	// 判断防火墙是否处于关闭状态
	if(rules.open_status == 0){
		// 关闭状态 放行数据包
		return NF_ACCEPT;
	}

	// 检查防火墙是否在工作时间段内 
	if(rules.settime_status == 1){
		ktime_t current_time = get_seconds();
		printk("current date： %ld  Start date： %ld  end date： %ld\n", current_time, rules.start_date, rules.end_date);
		if(current_time < rules.start_date || current_time >= rules.end_date){
			// 不在时间段内，直接放行数据包
			printk("Time: %s. The firewall is not in effect at the current time\n", time_buf);
			return NF_ACCEPT;
		}
	}

	// 状态检测
	if(rules.inp_status == 1){
		// 检查是否属于某个合法的链接
		if(check_conn(skb)){
			// printk("check passed\n");
			return NF_ACCEPT;  
		}
	}

	// 基于源IP地址访问控制
	if(rules.sip_status == 1){
		int sip_number;
		for (sip_number=0; sip_number < rules.sipNum; sip_number++){
			// 如果源地址 在被禁用地址(黑名单)内， 返回， 丢弃数据包
			if(rules.ban_sip[sip_number] == iph->saddr){
				printk("Time: %s. Request is deny. \nSrc IP: %s\tDest IP: %s\n\n", time_buf, src_ip_str, dst_ip_str);
				add_log(skb, "Source IP");
				return NF_DROP;
			}
		}
	}

	// 基于目的IP地址访问控制
	if(rules.dip_status == 1){
		int dip_number;
		for(dip_number=0; dip_number < rules.dipNum; dip_number++){
			if(rules.ban_dip[dip_number] == iph->daddr){
				printk("Time: %s. Request is deny. \nSrc IP: %s\tDest IP: %s\n\n", time_buf, src_ip_str, dst_ip_str);
				add_log(skb, "Destination IP");
				return NF_DROP;
			}
		}
	}

	// 基于源端口的访问控制
	if(rules.sport_status == 1){
		int sport_number;
		unsigned short sport, dport;
		printk("before  get port \n");
		get_port(skb, &sport, &dport);
		printk("get port ! S port : %hu\n", sport);
		for(sport_number = 0; sport_number < rules.sportNum; sport_number++){
			if(sport == rules.ban_sport[sport_number]){
				printk("Time: %s. Request is deny. \nSrc IP: %s\tDest IP: %s\tSrc port: %hu\n\n", time_buf, src_ip_str, dst_ip_str, sport);
				add_log(skb, "Source port");
				return NF_DROP;
			}
		}
	}

	// 基于目的端口的访问控制
	if(rules.dport_status == 1){
		int dport_number;
		unsigned short sport, dport;
		get_port(skb, &sport, &dport);
		for(dport_number = 0; dport_number < rules.dportNum; dport_number++){
			if(dport == rules.ban_sport[dport_number]){
				printk("Time: %s. Request is deny. \nSrc IP: %s\tDest IP: %s\tDest port: %hu\n\n", time_buf, src_ip_str, dst_ip_str, dport);
				add_log(skb, "Destination port");
				return NF_DROP;
			}
		}
	}

	// PING功能的控制
	if(rules.ping_status == 1){
		if(is_PING(iph)){
			printk("Time: %s. PING from %s is deny\n", time_buf, src_ip_str);
			add_log(skb, "PING Disabled");
			return NF_DROP;
		}
	}

	// HTTP/HTTPS 功能的控制
	if(rules.http_status == 1){
		tcph = tcp_hdr(skb);
		if(is_HTTP(iph, tcph)){
			printk("Time: %s. HTTP/HTTPS request is deny\n", time_buf);
			add_log(skb, "HTTP/HTTPS Disabled");
			return NF_DROP;				
		}
	}

	// Telnet 功能控制
	if(rules.telnet_status == 1){
		tcph = tcp_hdr(skb);
		if( is_TELNET(iph, tcph)){
			printk("Time: %s. Telnet request is deny\n", time_buf);
			add_log(skb, "Telnet Disabled");
			return NF_DROP;
		}
	}

	// 基于协议类型的访问控制
	if(rules.protocol_status == 1){

		switch (iph->protocol)
		{
			case IPPROTO_TCP:
				if(rules.protocol_type[0]){
					printk("Time: %s. TCP request is deny\n", time_buf);
					add_log(skb, "Protocol");
					return NF_DROP;
				}
				break;
			case IPPROTO_UDP:
				if (rules.protocol_type[1]){
					printk("Time: %s. UDP request is deny\n", time_buf);
					add_log(skb, "Protocol");
					return NF_DROP;
				}
				break;
			case IPPROTO_ICMP:
				if(rules.protocol_type[2]){
					printk("Time: %s. ICMP request is deny\n", time_buf);
					add_log(skb, "Protocol");
					return NF_DROP;
				}
				break;
			default:
				break;
		}
	}

	// 关闭连接功能的控制
	if (rules.close_status == 1){
		ktime_t current_time = get_seconds();

		if(rules.start_time <= current_time && rules.end_time >= current_time){
			printk("Time: %s. Request is deny. \nSrc IP: %s\tDest IP: %s\n\n", time_buf, src_ip_str, dst_ip_str);
			add_log(skb, "Close all connections");
			return NF_DROP;
		}
	}

	if(rules.inp_status == 1){
		update_hashTable(skb);
	}

    printk("Packet captured!\n");
    return NF_ACCEPT;
}

int hookSockoptSet(struct sock* sock, int cmd, sockptr_t user, unsigned int len){
    int ret;

    // 将用户空间的rules 拷贝到内核空间， 用recv接收，保证用户和内核空间过滤规则一致
    ret = copy_from_sockptr(&recv, user, sizeof(recv));

    switch(cmd){
        case OPENSTATE:         // 改变防火墙的开启状态
            rules.open_status = recv.open_status;
            break;
        case INPSTATE:
            rules.inp_status = recv.inp_status;
            rules.connNum = recv.connNum;
            memcpy(rules.connNode, recv.connNode, sizeof(rules.connNode));  
			// release_list(&connHead, &connEnd);
			// memset(&hashTable, 0, sizeof(hashTable));
            break;
        case NATSTATE:            // 改变防火墙NAT功能开启状态
			rules.nat_status = recv.nat_status;
			rules.natNum = recv.natNum;
			memcpy(rules.natTable, recv.natTable, sizeof(rules.natTable));  
			break;
		case SETTIME:             // 改变防火墙开启时间段
			rules.settime_status = recv.settime_status;
			rules.start_date = recv.start_date;
			rules.end_date = recv.end_date;
			break;
		case BANSIP:              // 基于源IP地址的访问控制
			rules.sip_status = recv.sip_status;
			rules.sipNum = recv.sipNum;
			memcpy(rules.ban_sip, recv.ban_sip, sizeof(rules.ban_sip));  
			break;
		case BANDIP:              // 基于目的IP地址的访问控制
			rules.dip_status = recv.dip_status;
			rules.dipNum = recv.dipNum;
			memcpy(rules.ban_dip, recv.ban_dip, sizeof(rules.ban_dip));  
			break;
		case BANSPORT:            // 基于源头端口的访问控制 
			rules.sport_status = recv.sport_status;
			rules.sportNum = recv.sportNum;    
			memcpy(rules.ban_sport, recv.ban_sport, sizeof(rules.ban_sport));  
			break;
		case BANDPORT:            // 基于目的端口的访问控制 
			rules.dport_status = recv.dport_status;
			rules.dportNum = recv.dportNum;   
			memcpy(rules.ban_dport, recv.ban_dport, sizeof(rules.ban_dport)); 
			break;
		case BANPROTOCOL:         // 基于协议类型的访问控制 
			rules.protocol_status = recv.protocol_status;
			memcpy(rules.protocol_type, recv.protocol_type, sizeof(rules.protocol_type)); 
			break;
		case BANMAC:              // 基于MAC地址的访问控制
			rules.mac_status = recv.mac_status;
			rules.macNum = recv.macNum;
			memcpy(rules.ban_mac, recv.ban_mac, sizeof(rules.ban_mac)); 
			break;
		case BANCOMBIN:           // 基于用户自定义策略的访问控制
			rules.combin_status = recv.combin_status;
			rules.combineNum = recv.combineNum; 
			memcpy(rules.ban_combin, recv.ban_combin, sizeof(rules.ban_combin)); 
			break;
		case BANALL:              // 关闭所有连接功能的控制
			rules.close_status = recv.close_status;
			rules.start_time = recv.start_time;
			rules.end_time = recv.end_time;
			break;
		case BANPING:             // PING功能的控制 
			rules.ping_status = recv.ping_status;
			break;
		case BANHTTP:             // HTTP/HTTPS功能的控制
			rules.http_status = recv.http_status;
			break;
		case BANTELNET:           // TELNET功能的控制
			rules.telnet_status = recv.telnet_status;
			break;
		case RESTORE:             // 恢复默认设置功能的控制
			memset(&rules, 0, sizeof(rules));	
			// release_list(&connHead, &connEnd);
			// memset(&hashTable, 0, sizeof(hashTable));
			rules.open_status = 1;
			break;
		default:
			break;
    }
    if (ret != 0)
	{
		ret = -EINVAL;
		printk("Error copying from user space to kernel space");
	}

    return ret;
}

int hookSockoptGet(struct sock* sock, int cmd, void __user* user, int* len){
    int ret;

    // if(cmd == CONNGET){
    //     int i = 0;
    //     Connection *p = connHead.next;
    // }

    // 将内核空间的rules 拷贝到用户空间 保证用户和内核空间过滤规则一致
    ret = copy_to_user(user, &rules, sizeof(rules));
    if(ret != 0){
        ret = -EINVAL;
        printk("Error copying from ketnel space to user space");
    }
    return ret;
}

int myfirewall_init(void) {
    //   初始化 后台周期性任务

    rules.open_status = 1;          // 初始化防火墙为开启状态
	rules.inp_status = 0;           // 初始化防火墙状态检测功能为关闭
	rules.nat_status = 0;           // 初始化防火墙NAT功能为关闭
	rules.sip_status = 0;           // 初始化基于源IP访问控制功能为关闭 
	rules.dip_status = 0;           // 初始化基于目的IP访问控制功能为关闭 
	rules.sport_status = 0;         // 初始化基于源端口访问控制功能为关闭  
	rules.dport_status = 0;         // 初始化基于目的端口访问控制功能为关闭 
	rules.settime_status = 0;	    // 初始化防火墙时间段功能为关闭
	rules.ping_status = 1;          // 初始化不封禁PING功能
	rules.http_status = 0;          // 初始化不封禁HTTP/HTTPS功能
	rules.telnet_status = 0;        // 初始化不封禁TELNET功能
	rules.protocol_status = 0;      // 初始化基于协议类型访问控制功能为关闭 
	rules.mac_status = 0;		    // 初始化基于MAC地址访问控制功能为关闭 
	rules.close_status = 0;         // 初始化开启所以连接
	rules.combin_status = 0; 	    // 初始化基于用户自定义策略访问控制功能为关闭

    // 设置钩子函数
    nfhoLocalIn.hook = hookLocalIn;
    nfhoLocalIn.pf = PF_INET;
    nfhoLocalIn.hooknum = NF_INET_LOCAL_IN;
    nfhoLocalIn.priority = NF_IP_PRI_FIRST;
    nf_register_net_hook(&init_net, &nfhoLocalIn);

    nfhoSockopt.pf = PF_INET;
    nfhoSockopt.set_optmax = SOE_MAX;
    nfhoSockopt.set_optmin = SOE_MIN;
    nfhoSockopt.set = hookSockoptSet;
    nfhoSockopt.get_optmin = SOE_MIN;
    nfhoSockopt.get_optmax = SOE_MAX;
    // nfhoSockopt.get = hookSockoptGet;

    nf_register_sockopt(&nfhoSockopt);

    printk("Firewall kernel module loaded successfully!\n");
    return 0;
}

void myfirewall_exit(void) {

    // 注销钩子
    nf_unregister_net_hook(&init_net, &nfhoLocalIn);

    printk("Firewall kernel module unloaded successfully!\n");
}

module_init(myfirewall_init);
module_exit(myfirewall_exit);

MODULE_LICENSE("GPL");
MODULE_AUTHOR("GY");
MODULE_DESCRIPTION("Netfilter Firewall v2.0");